package edu.buaa.energy.server;

import com.alibaba.fastjson.JSON;
import com.google.common.base.Preconditions;
import edu.buaa.common.RuntimeEnv;
import edu.buaa.common.server.DBSocketServer;
import edu.buaa.common.transaction.AbstractTransaction;
import edu.buaa.common.transaction.TransactionFailedException;
import edu.buaa.energy.EnergyDataTxGenerator;
import edu.buaa.energy.transaction.*;
import edu.buaa.utils.Helper;
import edu.buaa.utils.Pair;
import edu.buaa.utils.Triple;
import org.neo4j.graphdb.*;
import org.neo4j.graphdb.schema.IndexDefinition;
import org.neo4j.graphdb.schema.Schema;

import java.io.File;
import java.io.IOException;
import java.util.*;

public class Neo4jServer2 extends DBSocketServer.ReqExecutor {

    private final HashMap<Integer, Long> idMap2Long = new HashMap<>();

    public static void main(String[] args) {
        DBSocketServer server = new DBSocketServer(dbDir(), new Neo4jServer2(), 9866, "true".equals(System.getenv("ALLOW_TX_FAILURE")));
        RuntimeEnv env = RuntimeEnv.getCurrentEnv();
        String serverCodeVersion = env.name() + "." + Helper.codeGitVersion();
        System.out.println("server code version:" + serverCodeVersion);
        try {
            server.start();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }


    private static File dbDir() {
        String path = Helper.mustEnv("DB_PATH");
        Preconditions.checkNotNull(path, "need arg: DB_PATH");
        File dbDir = new File(path);
        if (!dbDir.exists()) {
            if (dbDir.mkdirs()) return dbDir;
            else throw new IllegalArgumentException("invalid dbDir");
        } else if (!dbDir.isDirectory()) {
            throw new IllegalArgumentException("invalid dbDir");
        }
        return dbDir;
    }

    private enum Edge implements RelationshipType {
        REACH_TO, PROP;
    }

    private Label label(String name){
        return DynamicLabel.label(name);
    }

    @Override
    public void setDB(GraphDatabaseService db) {
        this.db = db;
        try (Transaction tx = db.beginTx()) {
            Schema schema = db.schema();
            boolean hasTimeIndex = false;
            for(IndexDefinition timeIndex : schema.getIndexes()){
                hasTimeIndex = true;
            }
            if (!hasTimeIndex) {
                schema.indexFor(label("NODE")).on("id").create();
                System.out.println("Build the NODE index.");
                schema.indexFor(label("TPNODE")).on("st_time").create(); //.on("en_time")
                System.out.println("Build the TIME index.");
            }else{
                System.out.println("Has the time index already.");
            }
            tx.success();
        }
        buildIdMap(db);
    }

    private void buildIdMap(GraphDatabaseService db) {
        try (Transaction tx = db.beginTx();
             ResourceIterator<Node> allDataNodes = db.findNodes(label("NODE"))) {
            while (allDataNodes.hasNext()) {
                Node n = allDataNodes.next();
                Integer idInt = (Integer) n.getProperty("id");
                idMap2Long.put(idInt, n.getId());
            }
            tx.success();
        }
    }

    @Override
    protected AbstractTransaction.Result execute(String line) throws RuntimeException {
        try {
            AbstractTransaction tx = JSON.parseObject(line, AbstractTransaction.class);
            switch (tx.getTxType()) {
                case tx_import_static_data:
                    return execute((ImportStaticDataTx) tx);
                case tx_import_temporal_data:
                    return execute((ImportTemporalDataTx) tx);
                case tx_update_temporal_data:
                    return execute((UpdateTemporalDataTx) tx);
                case tx_query_snapshot:
                    return execute((SnapshotQueryTx) tx);
                case tx_query_snapshot_aggr_max:
                    return execute((SnapshotAggrMaxTx) tx);
                case tx_query_snapshot_aggr_duration:
                    return execute((SnapshotAggrDurationTx) tx);
                case tx_query_entity_history:
                    return execute((EntityHistoryTx) tx);
                case tx_query_road_by_temporal_condition:
                    return execute((EntityTemporalConditionTx) tx);
                default:
                    throw new UnsupportedOperationException();
            }
        }catch (NotFoundException e){
            throw new TransactionFailedException(e);
        }
    }

    private AbstractTransaction.Result execute(ImportStaticDataTx tx) {
        Label crossLabel = label("NODE");
        try (Transaction t = db.beginTx()) {
            for (ImportStaticDataTx.EnergyNode node : tx.getNodes()) {
                Node n = db.createNode(crossLabel);
                n.setProperty("id", node.getRawId());
                n.setProperty("name", node.getName());
                n.setProperty("country", node.getCountry());
                n.setProperty("voltage", node.getVoltage());
                n.setProperty("latitude", node.getLatitude());
                n.setProperty("longitude", node.getLongitude());
                idMap2Long.put(node.getRawId(), n.getId());
            }
            for (ImportStaticDataTx.EnergyRel rel : tx.getRels()) {
                Node s = db.getNodeById(idMap2Long.get(rel.getFromRawId()));
                Node e = db.getNodeById(idMap2Long.get(rel.getToRawId()));
                Relationship r = s.createRelationshipTo(e, Edge.REACH_TO);
                r.setProperty("x", rel.getX());
                r.setProperty("y", rel.getY());
                r.setProperty("limit", rel.getLimit());
                r.setProperty("length", rel.getLength());
            }
            t.success();
        }
        return new AbstractTransaction.Result();
    }

    private void setProperties(List<Float> values, Node r){
        for(int i=0; i< EnergyDataTxGenerator.properties.length; i++) {
            String propName = EnergyDataTxGenerator.properties[i];
            Float val = values.get(i);
            r.setProperty(propName, val);
        }
    }

    private Node createTemporalPropNode(int id, int startTime, int endTime) {
        Node ret = db.createNode(label("TPNODE"));
        ret.setProperty("node_id", id);
        ret.setProperty("st_time", startTime);
        ret.setProperty("en_time", endTime);
        return ret;
    }

    private AbstractTransaction.Result execute(ImportTemporalDataTx tx) {
        try (Transaction t = db.beginTx()) {
            for (ImportTemporalDataTx.EnergyStatus s : tx.getData()) {
                Long id = idMap2Long.get(s.getNodeId());
                assert id != null;
                Node n = db.getNodeById(id);
                int time = s.getTime();
                Node tp = createTemporalPropNode(s.getNodeId(), time, time+3599);
                setProperties(s.getStatus(), tp);
                n.createRelationshipTo(tp, Edge.PROP);
            }
            t.success();
        }
        return new AbstractTransaction.Result();
    }


    private void copyProperty(Node from, Node to){
        for(String propName : EnergyDataTxGenerator.properties) {
            to.setProperty(propName, from.getProperty(propName));
        }
    }

    private AbstractTransaction.Result execute(UpdateTemporalDataTx tx) {
        int tBegin = tx.getStartTime();
        int tFinish = tx.getEndTime();
        try (Transaction t = db.beginTx()) {
            // Do not consider that the old value equals to the new value
            Node node = db.getNodeById(idMap2Long.get(tx.getEntityId()));
            HashSet<Node> willBeDeleted = new HashSet<>();
            for(Relationship r: node.getRelationships(Edge.PROP)){
                Node tp = r.getOtherNode(node);
                int st = (int) tp.getProperty("st_time");
                int en = (int) tp.getProperty("en_time");
                if (tFinish >= st && tBegin <= en) { // has overlap
                    if (st < tBegin && tFinish < en) { // should split node.
                        Node theOtherPart = db.createNode(label("TPNODE"));
                        tp.setProperty("en_time", tBegin-1);
                        copyProperty(tp, theOtherPart);
                        theOtherPart.setProperty("st_time", tFinish+1);
                        theOtherPart.setProperty("en_time", en);
                        theOtherPart.setProperty("node_id", tx.getEntityId());
                        node.createRelationshipTo(theOtherPart, Edge.PROP);
                    }else if(tBegin <= st && en <= tFinish){ // should delete
                        willBeDeleted.add(tp);
                    }else if(tBegin <= st && tFinish < en ){ // change start
                        tp.setProperty("st_time", tx.getEndTime() + 1 );
                    }else if(st < tBegin && en <= tFinish){ // change end
                        tp.setProperty("en_time", tx.getStartTime() - 1 );
                    }
                }
            }
            for(Node toDel : willBeDeleted){
                for(Relationship r : toDel.getRelationships()){
                    r.delete();
                }
                toDel.delete();
            }
            Node insertedTp = db.createNode(label("TPNODE"));
            node.createRelationshipTo(insertedTp, Edge.PROP);
            insertedTp.setProperty("node_id", tx.getEntityId());
            insertedTp.setProperty("st_time", tx.getStartTime());
            insertedTp.setProperty("en_time", tx.getEndTime());
            setProperties(tx.getValues(), insertedTp);
            t.success();
        }
        return new AbstractTransaction.Result();
    }

    private AbstractTransaction.Result execute(SnapshotQueryTx tx) {
        try (Transaction t = db.beginTx()) {
            List<Pair<Integer, Float>> res = new ArrayList<>();
            int time = tx.getTimestamp();
            try(ResourceIterator<Node> it = db.findNodes(label("TPNODE"))) {
                while(it.hasNext()) {
                    Node tp = it.next();
                    int startTime = (int) tp.getProperty("st_time");
                    int endTime = (int) tp.getProperty("en_time");
                    if (startTime <= time && time <= endTime) {
                        Integer id = null;
                        for(Relationship r : tp.getRelationships()){
                            Node main = r.getOtherNode(tp);
                            id = (Integer) main.getProperty("id");
                        }
                        assert id!=null;
                        float v = (Float) tp.getProperty(tx.getPropertyName());
                        res.add(Pair.of(id, v));
                    }
                }
            }
            t.success();
            SnapshotQueryTx.Result result = new SnapshotQueryTx.Result();
            result.setStatus(res);
            return result;
        }
    }

    private AbstractTransaction.Result execute(SnapshotAggrMaxTx tx) {
        try (Transaction t = db.beginTx()) {
            List<Pair<Integer, Float>> res = new ArrayList<>();
            HashMap<Integer, Float> tmp = new HashMap<>();
            int st = tx.getT0();
            int en = tx.getT1();
            try(ResourceIterator<Node> it = db.findNodes(label("TPNODE"))) {
                while(it.hasNext()) {
                    Node tp = it.next();
                    int startTime = (int) tp.getProperty("st_time");
                    int endTime = (int) tp.getProperty("en_time");
                    if (startTime <= en && st <= endTime) {
                        Integer id = null;
                        for(Relationship r : tp.getRelationships()){
                            Node main = r.getOtherNode(tp);
                            id = (Integer) main.getProperty("id");
                        }
                        assert id!=null;
                        float v = (Float) tp.getProperty(tx.getP());
                        tmp.merge(id, v, Float::max);
                    }
                }
            }
            t.success();
            tmp.forEach((k, v) -> res.add(Pair.of(k, v)));
            SnapshotAggrMaxTx.Result result = new SnapshotAggrMaxTx.Result();
            result.setPropMaxValue(res);
            return result;
        }
    }

    private AbstractTransaction.Result execute(SnapshotAggrDurationTx tx) {
        try (Transaction t = db.beginTx()) {
            List<Triple<Integer, Float, Integer>> res = new ArrayList<>();
            // id, property_value, duration
            HashMap<Integer, HashMap<Float, Integer>> buffer = new HashMap<>();
            int t0 = tx.getT0();
            int t1 = tx.getT1();
            try(ResourceIterator<Node> it = db.findNodes(label("TPNODE"))) {
                while(it.hasNext()) {
                    Node tp = it.next();
                    int st = (int) tp.getProperty("st_time");
                    int en = (int) tp.getProperty("en_time");
                    if (t1 >= st && t0 <= en) {
                        int left = Math.max(t0, st);
                        int right = Math.min(t1, en);
                        int duration = right - left + 1;
                        Integer name = null;
                        for(Relationship r : tp.getRelationships()){
                            Node main = r.getOtherNode(tp);
                            name = (Integer) main.getProperty("id");
                        }
                        assert name!=null;
                        float value = (float) tp.getProperty(tx.getP());
                        Float grp = tx.getIntStartTreeSet().floor(value);
                        if(grp==null) continue;
                        HashMap<Float, Integer> tmp = buffer.get(name);
                        if (tmp != null) {
                            tmp.merge(grp, duration, Integer::sum);
                        } else {
                            tmp = new HashMap<>();
                            tmp.put(grp, duration);
                            buffer.put(name, tmp);
                        }
                    }
                }
            }
            t.success();
            buffer.forEach((key, value) -> value.forEach((k, v) -> res.add(Triple.of(key, k, v))));
            SnapshotAggrDurationTx.Result result = new SnapshotAggrDurationTx.Result();
            result.setStatusDuration(res);
            return result;
        }
    }

    private AbstractTransaction.Result execute(EntityHistoryTx tx) {
        try (Transaction t = db.beginTx()) {
            List<Triple<Integer, Integer, Float>> res = new ArrayList<>();
            int t0 = tx.getBeginTime();
            int t1 = tx.getEndTime();
            try(ResourceIterator<Node> it = db.findNodes(label("NODE"), "id", tx.getEntity())) {
                if(it.hasNext()) {
                    Node node = it.next();
                    for(Relationship r : node.getRelationships(Edge.PROP)){
                        Node tp = r.getOtherNode(node);
                        int st = (int) tp.getProperty("st_time");
                        int en = (int) tp.getProperty("en_time");
                        float value = (float) tp.getProperty(tx.getProp());
                        if (t1 >= st && t0 <= en) {
                            int left = Math.max(t0, st);
                            int right = Math.min(t1, en);
                            res.add(Triple.of(left, right, value));
                        }
                    }
                }
            }
            t.success();
            EntityHistoryTx.Result result = new EntityHistoryTx.Result();
            result.setHistory(res);
            return result;
        }
    }

    private AbstractTransaction.Result execute(EntityTemporalConditionTx tx) {
        int st = tx.getT0();
        int en = tx.getT1();
        float vMin = tx.getVMin();
        float vMax = tx.getVMax();
        HashSet<Integer> res = new HashSet<>();
        try (Transaction t = db.beginTx();
             ResourceIterator<Node> it = db.findNodes(label("TPNODE"))) {
            while(it.hasNext()) {
                Node tp = it.next();
                int startTime = (int) tp.getProperty("st_time");
                int endTime = (int) tp.getProperty("en_time");
                float v = (Float) tp.getProperty(tx.getP());
                if (startTime <= en && st <= endTime && vMin<=v && v<=vMax) {
                    Integer name = null;
                    for(Relationship r : tp.getRelationships()){
                        Node main = r.getOtherNode(tp);
                        name = (Integer) main.getProperty("id");
                    }
                    assert name!=null;
                    res.add(name);
                }
            }
            t.success();
        }
        EntityTemporalConditionTx.Result result = new EntityTemporalConditionTx.Result();
        result.setEntities(new ArrayList<>(res));
        return result;
    }

}
